import tempfile
import shutil
from contextlib import contextmanager

from paramiko import SSHClient, Channel, AutoAddPolicy

from tinyops.library.task import TaskLogger, TaskEnv
from tinyops.library.config import ServerConfig

from .helper import ExportHelper


class SSHExecutor:
    def __init__(self, server: ServerConfig, env: TaskEnv, timeout: int, logger: TaskLogger):
        self.name = server.name
        self._log = logger.target(self.name).log
        self.server = server
        self.env = env if env is not None else TaskEnv()
        self.timeout = timeout
        self.logger = logger
        self.buffer_size = -1
        self.client = SSHClient()
        self.client.set_missing_host_key_policy(AutoAddPolicy())
        self._export_helper = ExportHelper(
            env=self.env,
            logger=logger.target(self.name),
            handler=self._export_handler,
        )
        self._inited = False

    def init(self):
        if self._inited:
            return
        self._inited = True
        self.client.connect(
            self.server.host,
            self.server.port,
            username=self.server.user,
            password=self.server.password or None,
            key_filename=self.server.private_key or None,
            timeout=self.timeout
        )
        # keep connections alive by send ignore packet every interval seconds
        self.client.get_transport().set_keepalive(30)
        self._log('server {}@{}:{} connected'.format(
            self.server.user,
            self.server.host,
            self.server.port,
        ))

    def close(self):
        self.client.close()
        self._log('server {}@{}:{} disconnected'.format(
            self.server.user,
            self.server.host,
            self.server.port,
        ))

    @staticmethod
    def _exec_command_env(chan: Channel, script: str, env: dict):
        """
        Execute export command to set environment.

        The environment argument in Channel.exec_command may not work.
        The server may reject this request depending on its ``AcceptEnv``
        setting; such rejections will fail silently.
        See also: Warning in paramiko.Channel.set_environment_variable
        """
        commands = [f'export {k}={v};' for k, v in env.items()]
        full_script = '\n'.join(commands) + '\n' + script
        chan.exec_command(full_script)

    @contextmanager
    def _exec_command(self, script: str, env: dict = None):
        chan = self.client.get_transport().open_session(timeout=self.timeout)
        try:
            chan.get_pty()
            chan.settimeout(self.timeout)
            if env:
                self._exec_command_env(chan, script, env=env)
            else:
                chan.exec_command(script)
            stdout = chan.makefile("r", self.buffer_size)
            yield chan, stdout
        finally:
            chan.close()
            stdout.close()

    def _export_handler(self, export_file):
        local_path = tempfile.mktemp("tinyops_temp_export_file")
        try:
            with self.client.open_sftp() as sftp:
                sftp.get(export_file, local_path)
        except IOError as ex:
            self._log(f'EXPORT failed, can not transfer export file: {ex}')
            return None
        with open(local_path) as f:
            content = f.read()
        shutil.rmtree(local_path, ignore_errors=True)
        clean_command = f'rm -f {export_file}'
        with self._exec_command(clean_command) as (chan, output):
            exit_code = str(chan.recv_exit_status())
            if exit_code != '0':
                self._log(f'Warning: delete temp file failed: {export_file}')
        return content

    def do_run(self, script: str, exports: list = None):
        self._log(script, sep='>')
        script = 'set -e;\n' + script
        script = self._export_helper.before_run(script, exports)
        with self._exec_command(script, env=self.env.value()) as (chan, stdout):
            outputs = []
            for line in stdout:
                outputs.append(line)
                self._log(line)
            exit_code = str(chan.recv_exit_status())
        self._log(exit_code=exit_code)
        if not self._export_helper.after_run(outputs):
            return False
        return exit_code == '0'

    def do_get(self, remote_path, local_path):
        self._log(f'## GET {remote_path} {local_path}', sep='>')
        with self.client.open_sftp() as sftp:
            sftp.get(remote_path, local_path)
        return True

    def do_put(self, local_path, remote_path):
        self._log(f'## PUT {local_path} {remote_path}', sep='>')
        with self.client.open_sftp() as sftp:
            sftp.put(local_path, remote_path)
        return True
